# """公司对象操作方法包
# @Project: Rebar
# @File: company_mapper.py
# @Author: hzc
# @Date: 2023/3/27 20:49
# """
# from typing import List
# from sqlalchemy.exc import IntegrityError
# # from model.db import rebar, TbCompany
# from model.db import rebar, TbCompanies
# from model.params import CompanyInfoParam
# from utils import DataNotExist, AllException
#
#
from pymysql import IntegrityError
from model.params import CompanyInfoParam
from model.db import TbCompanies, reservoir_sensitivity
from util import DataNotExist


def get_company(company_name: str) -> TbCompanies:
    """
    根据公司名获取公司信息
    :param company_name:
    :return:
    """
    return company_is_exist(company_name)


def get_company_Bycode(company_code: str) -> TbCompanies:
    """
    根据公司名获取公司信息
    :param company_code:

    :return:
    """
    return companyCode_is_exist(company_code)


def get_company_list():
    """
    根据公司名获取公司信息
    :param company_name:
    :return:
    """
    return reservoir_sensitivity.query(TbCompanies).all()


#
def update_company_info(add, delete, update):
    """
    更新开发单元
    """
    try:
        if len(add) != 0:
            for item in add:
                # company = TbCompanies(
                #     company_name=item['company_name'],
                #     company_code=item['company_code'],
                #     company_representative=item['company_representative'],
                #     remark=item['company_representative']
                # )
                company = TbCompanies(
                    **item
                )
                reservoir_sensitivity.add(company)
                reservoir_sensitivity.commit()
                item['company_id'] = company.company_id
        if len(delete) != 0:
            for item in delete:
                reservoir_sensitivity.query(TbCompanies).filter(TbCompanies.company_id == item['company_id']) \
                    .delete(synchronize_session=False)

        if len(update) != 0:
            for item in update:
                reservoir_sensitivity.query(TbCompanies).filter(TbCompanies.company_id == item['company_id']) \
                    .update(item)
        reservoir_sensitivity.commit()
        reservoir_sensitivity.flush()
        return add
    except Exception as e:
        reservoir_sensitivity.rollback()
        raise e
def add_company(company_info: CompanyInfoParam) -> TbCompanies:
    """
    添加公司
    """
    if reservoir_sensitivity.query(TbCompanies).filter(TbCompanies.company_id == 1).first():
        return
    try:
        company = TbCompanies(
            company_name=company_info.company_name,
            company_code=company_info.company_code,
            company_representative=company_info.company_representative,
            remark=company_info.remark
        )
        reservoir_sensitivity.add(company)
        reservoir_sensitivity.commit()
        reservoir_sensitivity.refresh(company)
        return company
    except IntegrityError as e:
        print("company_params.py", e, e.__traceback__.tb_lineno, sep="\n")
        reservoir_sensitivity.rollback()
        raise e


#
# def update_company(company_id: int, company_info: CompanyInfoParam) -> TbCompany:
#     """
#     修改公司信息
#     """
#     company = company_is_exist(company_id)
#     try:
#         for key, value in dict(company_info).items():
#             setattr(company, key, value)
#         rebar.commit()
#         rebar.refresh(company)
#         return company
#     except Exception as e:
#         print("company_params.py", e, e.__traceback__.tb_lineno)
#         rebar.rollback()
#         raise AllException("修改公司信息失败")
#
#
# def delete_company(company_id: int) -> bool:
#     """
#     删除公司
#     """
#     company = company_is_exist(company_id)
#     try:
#         company.is_delete = 1
#         rebar.commit()
#         rebar.refresh(company)
#         return True
#     except Exception as e:
#         print("company_params.py", e, e.__traceback__.tb_lineno)
#         rebar.rollback()
#         raise AllException("删除公司信息失败")
#
#
# def get_company_by_name(company_name) -> TbCompany:
#     company = rebar.query(TbCompany).filter_by(company_name=company_name).first()
#     if not company:
#         raise DataNotExist("公司")
#     return company
#
#
def company_is_exist(company_name) -> TbCompanies:
    company = reservoir_sensitivity.query(TbCompanies).filter_by(company_name=company_name).first()
    if company is None:
        raise DataNotExist("公司")
    return company


def companyCode_is_exist(company_code) -> TbCompanies:
    company = reservoir_sensitivity.query(TbCompanies).filter_by(company_code=company_code).first()
    if company is None:
        raise DataNotExist("公司")
    return company
